{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to use LLMs and SQL?\n",
    "\n",
    "## Comparision Between methods\n",
    "| Method | About | Pros | Cons |\n",
    "|----------|----------|----------|----------|\n",
    "| SQLCoder-7b    |          |          |          |\n",
    "| BentoML    |          |          |          |\n",
    "| Row 3    |          |          |          |\n",
    "\n",
    "## Basic Setup\n",
    "Most LLMs that do SQL generation need your SQL Schema and a connection to your database to run the query.\n",
    "\n",
    "### SQLCoder\n",
    "Comes in multiple varients, with the largest being 70b, 34b & 7b. \n",
    "\n",
    "Using the 7b-v2 we can build a small viable demo:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading the model:\n",
    "model_name = \"defog/sqlcoder-7b-2\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "if available_memory > 17e9:\n",
    "    # if you have atleast 16GB of GPU memory, run load the model in float16\n",
    "    model = AutoModelForCausalLM.from_pretrained(\n",
    "        model_name,\n",
    "        trust_remote_code=True,\n",
    "        torch_dtype=torch.float16,\n",
    "        device_map=\"auto\",\n",
    "        use_cache=True,\n",
    "    )\n",
    "else:\n",
    "    # else, load in 8 bits – this is a bit slower\n",
    "    model = AutoModelForCausalLM.from_pretrained(\n",
    "        model_name,\n",
    "        trust_remote_code=True,\n",
    "        # torch_dtype=torch.float16,\n",
    "        load_in_4bit=True,\n",
    "        device_map=\"auto\",\n",
    "        use_cache=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting the prompt correctly is essential:\n",
    "prompt = \"\"\"### Task\n",
    "Generate a SQL query to answer [QUESTION]{question}[/QUESTION]\n",
    "\n",
    "### Instructions\n",
    "- If you cannot answer the question with the available database schema, return 'I do not know'\n",
    "- Remember that revenue is price multiplied by quantity\n",
    "- Remember that cost is supply_price multiplied by quantity\n",
    "\n",
    "### Database Schema\n",
    "This query will run on a database whose schema is represented in this string:\n",
    "CREATE TABLE products (\n",
    "  product_id INTEGER PRIMARY KEY, -- Unique ID for each product\n",
    "  name VARCHAR(50), -- Name of the product\n",
    "  price DECIMAL(10,2), -- Price of each unit of the product\n",
    "  quantity INTEGER  -- Current quantity in stock\n",
    ");\n",
    "\n",
    "CREATE TABLE customers (\n",
    "   customer_id INTEGER PRIMARY KEY, -- Unique ID for each customer\n",
    "   name VARCHAR(50), -- Name of the customer\n",
    "   address VARCHAR(100) -- Mailing address of the customer\n",
    ");\n",
    "\n",
    "CREATE TABLE salespeople (\n",
    "  salesperson_id INTEGER PRIMARY KEY, -- Unique ID for each salesperson\n",
    "  name VARCHAR(50), -- Name of the salesperson\n",
    "  region VARCHAR(50) -- Geographic sales region\n",
    ");\n",
    "\n",
    "CREATE TABLE sales (\n",
    "  sale_id INTEGER PRIMARY KEY, -- Unique ID for each sale\n",
    "  product_id INTEGER, -- ID of product sold\n",
    "  customer_id INTEGER,  -- ID of customer who made purchase\n",
    "  salesperson_id INTEGER, -- ID of salesperson who made the sale\n",
    "  sale_date DATE, -- Date the sale occurred\n",
    "  quantity INTEGER -- Quantity of product sold\n",
    ");\n",
    "\n",
    "CREATE TABLE product_suppliers (\n",
    "  supplier_id INTEGER PRIMARY KEY, -- Unique ID for each supplier\n",
    "  product_id INTEGER, -- Product ID supplied\n",
    "  supply_price DECIMAL(10,2) -- Unit price charged by supplier\n",
    ");\n",
    "\n",
    "-- sales.product_id can be joined with products.product_id\n",
    "-- sales.customer_id can be joined with customers.customer_id\n",
    "-- sales.salesperson_id can be joined with salespeople.salesperson_id\n",
    "-- product_suppliers.product_id can be joined with products.product_id\n",
    "\n",
    "### Answer\n",
    "Given the database schema, here is the SQL query that answers [QUESTION]{question}[/QUESTION]\n",
    "[SQL]\n",
    "\"\"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generating the SQL Query:\n",
    "import sqlparse\n",
    "\n",
    "def generate_query(question):\n",
    "    updated_prompt = prompt.format(question=question)\n",
    "    inputs = tokenizer(updated_prompt, return_tensors=\"pt\").to(\"cuda\")\n",
    "    generated_ids = model.generate(\n",
    "        **inputs,\n",
    "        num_return_sequences=1,\n",
    "        eos_token_id=tokenizer.eos_token_id,\n",
    "        pad_token_id=tokenizer.eos_token_id,\n",
    "        max_new_tokens=400,\n",
    "        do_sample=False,\n",
    "        num_beams=1,\n",
    "    )\n",
    "    outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)\n",
    "\n",
    "    torch.cuda.empty_cache()\n",
    "    torch.cuda.synchronize()\n",
    "    # empty cache so that you do generate more results w/o memory crashing\n",
    "    # particularly important on Colab – memory management is much more straightforward\n",
    "    # when running on an inference service\n",
    "    return sqlparse.format(outputs[0].split(\"[SQL]\")[-1], reindent=True)\n",
    "\n",
    "question = \"What was our revenue by product in the new york region last month?\"\n",
    "generated_sql = generate_query(question)\n",
    "print(generated_sql)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## See Also\n",
    "- [Chat Bot](https://github.com/Exorust/LLM-Cookbook/blob/main/chatbot.md)\n",
    "\n",
    "## References\n",
    "\n",
    "This project uses code from the following source:\n",
    "- **SQLCoder**: Available at: [URL to the original source](https://github.com/defog-ai/sqlcoder?tab=readme-ov-file)\n",
    "- **SQLCoder Demo Code**: Available at: [URL to the original source](https://colab.research.google.com/drive/1z4rmOEiFkxkMiecAWeTUlPl0OmKgfEu7?usp=sharing#scrollTo=_mNT25UOfSMw)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
